"""
均线交叉策略

经典的双均线交叉策略：
- 短期均线上穿长期均线时买入（金叉）
- 短期均线下穿长期均线时卖出（死叉）

作者: AI Assistant
版本: 1.0.0
日期: 2025-01-06
"""

import pandas as pd
from backtest.strategy import StrategyBase


class MAStrategy(StrategyBase):
    """
    均线交叉策略
    
    使用两条移动平均线的交叉作为买卖信号。
    
    参数:
        ma_short: 短期均线周期，默认5
        ma_long: 长期均线周期，默认20
        
    信号规则:
        - 金叉（短期MA上穿长期MA）：买入
        - 死叉（短期MA下穿长期MA）：卖出
        - 其他情况：持有
        
    示例:
        >>> strategy = MAStrategy(ma_short=5, ma_long=20)
        >>> engine = BacktestEngine(strategy=strategy, data=data)
        >>> result = engine.run()
    """
    
    def __init__(self, ma_short: int = 5, ma_long: int = 20):
        """
        初始化均线策略
        
        参数:
            ma_short: 短期均线周期
            ma_long: 长期均线周期
        """
        # 参数验证
        if ma_short >= ma_long:
            raise ValueError("短期均线周期必须小于长期均线周期")
        
        if ma_short < 1 or ma_long < 1:
            raise ValueError("均线周期必须大于0")
        
        super().__init__(
            name=f"均线交叉策略(MA{ma_short}/MA{ma_long})",
            min_period=ma_long,  # 至少需要长期均线周期的数据
            ma_short=ma_short,
            ma_long=ma_long
        )
        
        self.ma_short = ma_short
        self.ma_long = ma_long
        
        # 用于记录上一个信号，避免重复交易
        self.last_signal = None
    
    def initialize(self, data: pd.DataFrame):
        """
        初始化策略，预计算均线
        
        参数:
            data: 历史数据
        """
        super().initialize(data)
        # 这里可以预计算一些指标，但为了回测的真实性，
        # 我们在generate_signal中实时计算
    
    def generate_signal(self, data: pd.DataFrame, current_bar: pd.Series) -> str:
        """
        生成交易信号
        
        参数:
            data: 当前及之前的所有数据
            current_bar: 当前K线数据
            
        返回:
            str: 'BUY', 'SELL', 或 'HOLD'
        """
        # 确保有足够的数据
        if len(data) < self.ma_long:
            return 'HOLD'
        
        # 计算短期和长期均线
        # 使用最近的数据计算，避免未来函数
        ma_short_series = data['收盘'].rolling(window=self.ma_short, min_periods=self.ma_short).mean()
        ma_long_series = data['收盘'].rolling(window=self.ma_long, min_periods=self.ma_long).mean()
        
        # 获取当前和前一天的均线值
        if len(ma_short_series) < 2 or len(ma_long_series) < 2:
            return 'HOLD'
        
        ma_short_current = ma_short_series.iloc[-1]
        ma_long_current = ma_long_series.iloc[-1]
        ma_short_prev = ma_short_series.iloc[-2]
        ma_long_prev = ma_long_series.iloc[-2]
        
        # 检测金叉（买入信号）
        # 前一天短期MA在长期MA下方，今天上穿
        if ma_short_prev <= ma_long_prev and ma_short_current > ma_long_current:
            if self.last_signal != 'BUY':
                self.last_signal = 'BUY'
                return 'BUY'
        
        # 检测死叉（卖出信号）
        # 前一天短期MA在长期MA上方，今天下穿
        elif ma_short_prev >= ma_long_prev and ma_short_current < ma_long_current:
            if self.last_signal != 'SELL':
                self.last_signal = 'SELL'
                return 'SELL'
        
        return 'HOLD'
    
    def get_indicator_values(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        获取策略指标值（用于绘图）
        
        参数:
            data: 历史数据
            
        返回:
            pd.DataFrame: 包含均线值的数据框
        """
        result = data.copy()
        result[f'MA{self.ma_short}'] = data['收盘'].rolling(
            window=self.ma_short,
            min_periods=self.ma_short
        ).mean()
        result[f'MA{self.ma_long}'] = data['收盘'].rolling(
            window=self.ma_long,
            min_periods=self.ma_long
        ).mean()
        return result

